#!/usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt

def show():
    plt.show()

def barplot(means, stds, colors, bar_width, \
            title, y_label, x_ticks, y_scale_min, y_scale_max, \
            print_values=False):
    nBarPerGroup = len(means)
    bar_width = bar_width / nBarPerGroup
    n = len(means[0])

    ind = np.arange(n)  # the x locations for the groups
    fig = plt.figure()
    ax = fig.add_subplot(111)

    rect_hdl = []
    lines = []
    for c in range(nBarPerGroup):
        tmp = ax.bar(ind, means[c], bar_width, color=colors[c], yerr=stds[c])
        rect_hdl.append(tmp)
        lines.append(tmp[0])
        ind = ind + bar_width

    # add some
    ax.set_ylabel(y_label)
    ax.set_title(title)
    ax.set_xticks(np.arange(n) + bar_width * nBarPerGroup / 2)
    ax.set_xticklabels(x_ticks)
    ax.set_xlim(xmin= -bar_width, max=ind + 2 * (bar_width))
    ax.set_ybound(lower=y_scale_min, upper=y_scale_max)

    ax.legend(lines, label)

    def autolabel(rects):
        for rect in rects:
            height = rect.get_height()
            ax.text(rect.get_x() + rect.get_width() / 2., 1.05 * height, '%d' % int(height),
                    ha='center', va='bottom')

    if print_values:
        for c in range(nBarPerGroup):
            autolabel(rect_hdl[c])

    return fig


def SNARC_plot(x, y, beta, intercept, title):
    fig = plt.figure()
    ax = fig.add_subplot(111)

    #line
    xx = np.array([x[0] - 0.5, x[-1] + 0.5])
    ax.plot(xx, beta * xx + intercept, linewidth=2, color=[.5, .5, .5])
    #dots
    ax.plot(x, y, 'o', color=[.2, .2, .2], markersize=8)
    ax.set_xlim(0, 10)
    ax.set_aspect(0.75 * np.diff(ax.get_xlim()) / np.diff(ax.get_ylim()))
    ax.set_xticks(x)
    ax.set_title(title)
    return fig



if __name__ == "__main__":
#DEMOS

    if (True):
        #BAR PLOT
        label = ["men", "Women"]
        means = [[20, 35, 30, 35, 27], [25, 32, 34, 20, 25] ]
        stds = [None, None]
        colors = [[0.1, 0.1, 0.1], [0.8, 0.8, 0.8]]

        b = barplot(means, stds, colors,
                bar_width=0.70, \
                title="Scores by group and gender",
                y_label="Scores",
                x_ticks=['G1', 'G2', 'G3', 'G4', 'G5'],
                y_scale_min=0,
                y_scale_max=40)

        ax = b.get_axes()[0]
        ax.set_aspect(0.75 * np.diff(ax.get_xlim()) / np.diff(ax.get_ylim()))
        #b.savefig('hello.eps', format='eps')

    if (True):
        #DEMO "SNARC PLOT" 
        x = [1, 2, 3, 4, 6, 7, 8, 9]
        y = [46.392857, 29.14285714, -21.64285714, 5.178571429, -29.46428571, -8.357142857, -43.71428571, -34.53571429]
        fig = SNARC_plot(x, y, beta= -9.1726, intercept=38.7381, title="Iranian Teenage Group 1")
        #fig.savefig("SNARC", format='eps')

    plt.show()

